…xport

transformer_sana_wm.py:
* License header switched to the "HuggingFace Team and SANA-WM Authors"
  style used by merged sana_video.
* Imports rewritten in stdlib -> third-party -> diffusers order; use
  diffusers `from ...utils import logging` instead of stdlib `logging`.
* Fix 9 `Optional[X]` annotations written as `X or None` (Python's `or`
  short-circuits and silently returns `X`).
* Fix two `assert (cond, msg)` tuple-asserts in PatchEmbedMS3D.forward
  that always pass (SyntaxWarning at import time).
* Remove duplicate `__all__` declarations (the second silently overwrote
  the first).
* Remove dead `reset_bn` (imports a nonexistent `packages.apps.utils`,
  would crash on call).
* Remove the duplicate `logger = logging.getLogger(__name__)` further
  down in the file.

transformer_sana_wm_kernels.py:
* License header normalized; collapse three duplicate triton/torch import
  blocks into one.

pipeline_sana_wm.py:
* License header normalized.
* `_decode_latents` now returns `(T, H, W, 3)` float in [0, 1], matching
  the diffusers convention used by `VideoProcessor`. Returning uint8
  silently broke `export_to_video`: it does `frame * 255` assuming float
  input, so uint8 overflows to `(-x) mod 256` and inverts colors.
* `__call__` converts to PIL/uint8 only when `output_type="pil"`.
* Intrinsics argument now accepts (4,), (F, 4), (3, 3), and (F, 3, 3)
  forms (auto-extracts fx, fy, cx, cy from a 3x3 K) and auto-trims to
  `num_frames` when a longer-than-needed trajectory is passed.
* Inline `retrieve_timesteps` with the standard `# Copied from
  diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps`
  marker, matching merged sana_video.
* Docstrings + EXAMPLE_DOC_STRING updated to reflect the new return type.

pipeline_output.py:
* Update `frames` field docstring to describe the new float [0, 1] return.

refiner.py, cam_utils.py, scripts/sana_wm/convert_sana_wm_to_diffusers.py:
* License headers normalized.

Docs:
* New `docs/source/en/api/pipelines/sana_wm.md` and
  `docs/source/en/api/models/sana_wm_transformer3d.md`, modeled on
  sana_video.md / sana_video_transformer3d.md, wired into
  `docs/source/en/_toctree.yml` under Models and Pipelines.

5s end-to-end smoke test (81 frames @ 16fps, 30 stage-1 steps + 3-step
LTX-2 refiner) passes on 1x H100 80GB with `enable_model_cpu_offload`.
Round-trip diff vs raw float frames is 2.06/255 mean (h264 lossy noise),
confirming the export_to_video fix.

…+ KV cache hooks)

The first cleanup pass only kept the legacy single-shot refiner path. That
path is what the model was *not* trained on — its docstring even says
"feeding the full sequence at once is out-of-distribution" — and its cost
is O(T^2) attention over the full latent volume, which made longer videos
unusable (~21 min per refiner step at 321 frames on an H100).

Port the chunk-causal AR mode from the upstream reference so the refiner
matches the training contract:

* `refine_latents` now defaults to `block_size=3, kv_max_frames=11`
  (the canonical AR recipe). Pass `block_size=None` to fall back to the
  legacy single-shot path.
* New `_refine_latents_ar` + `_RefinerChunkRunner` orchestrate the sliding
  window: pre-capture pre-RoPE sink K/V on `z_sana[:source_sink_frames]`
  at sigma=0, then for each `block_size`-frame chunk run a 3-step Euler
  with prefix `{sink_k_pre, sink_v, sink_pe, history_k, history_v}` and
  capture post-RoPE K/V to feed the next window. History is bounded to
  `kv_max_frames - source_sink_frames` so per-block compute is constant.
* New `_predict_x0_active_block` runs the transformer on the active block
  only (Q from active, K/V from prefix+active).
* New `_capture_block_kv` runs sigma=0 forward with a pre_rope/post_rope
  capture flag set on each `attn1`.
* New `_forward_video_only_with_rope` takes a pre-built RoPE so each block
  can use absolute frame positions in the source video.
* `_streaming_self_attention` extended with the `_kv_cache_capture`,
  `_tf_capture_kv`, `_tf_kv_prefix` hook contract that AR mode uses to
  inject and capture K/V on each block.
* New helpers: `_build_rotary_emb_for_absolute_positions`,
  `_set_kv_prefix_on_blocks`, `_clear_kv_prefix_on_blocks`,
  `_set_capture_flag_on_blocks`, `_collect_captured_kv_from_blocks`.
* `_encode_prompt` now also moves the Gemma-3 text encoder back to CPU
  after producing the embeds — otherwise it stays resident through the
  entire AR loop and gates how much GPU memory the refiner transformer
  has left.

Module-level docstring updated to document both modes; existing
single-shot path preserved verbatim.

…eemption)

The AR refiner is expensive (~3-5 min per block) and the refinement loop
ran end-to-end has no in-progress state to recover, so a SLURM preemption
mid-refinement loses all progress. With the canonical
``block_size=3, kv_max_frames=11`` setup, refining a 50s video is 34
blocks of work that has to make it through without preemption on a
backfill queue.

Add per-block atomic checkpointing:

* ``SanaWMLTX2Refiner.refine_latents(checkpoint_dir=Path)`` and
  ``_refine_latents_ar`` accept a directory. After each completed AR
  block, the AR loop writes ``checkpoint_dir/state.pt`` atomically
  (tmp + os.replace).
* The payload is ``{block_idx_done, n_blocks, sink_size, block_size,
  output_shape, output, runner_state}``. ``runner_state`` is a CPU snapshot
  of the runner's ``_sink_kv_pre``, ``_history_kv_post``,
  ``_history_frames`` and ``torch.Generator`` state.
* On entry, if ``state.pt`` exists with a compatible shape signature, the
  AR loop loads the persisted output tensor + runner state and resumes
  from ``block_idx_done + 1`` instead of recomputing from scratch.
* ``SanaWMPipeline.__call__(refiner_checkpoint_dir=...)`` plumbs the
  directory through to the refiner.

Checkpoint size: ~output_volume + sink_KV (~360MB for 50 layers) +
rolling history KV (~3-4GB at full capacity) — saved once per block,
total per-block save overhead ~10s on lustre.

Three CI checks were failing on the PR:

1. `check_code_quality` (43 ruff errors): mix of unused imports / import
   sorting / E731 lambdas (auto-fixable) plus a handful of F821 dead-code
   references inherited from the upstream research codebase (`xformers.*`
   inside `if _xformers_available:` blocks, an undefined `BlockHook` type
   annotation, two `x_sa`/`mlp_out` references in a block forward whose
   live assignment was already overridden by subclasses). Ran `ruff check
   --fix --unsafe-fixes` + `ruff format`, fixed the type annotation
   manually, and added targeted `# noqa: F821` markers on the conditionally
   unreachable lines.

2. `check_torch_dependencies`: `transformer_sana_wm.py` hard-imported
   `einops`, `fla`, `timm`, `termcolor`. The minimum-deps CI environment
   doesn't have them, and diffusers' lazy loader rewrites `ModuleNotFoundError`
   as `RuntimeError` so `test_pipeline_imports` blew up. Wrapped each of
   the four optional imports in a try/except shim — `rearrange`/
   `ShortConvolution`/`DropPath`/`Attention_`/`Mlp` become placeholders
   that raise a clear `ImportError` on construction, `colored` falls back
   to plain text. Class bodies that subclass these still parse at module
   load, so `import diffusers.models.transformers.transformer_sana_wm`
   succeeds anywhere. Same treatment for the kernels file's
   `from einops import rearrange, repeat`.

3. `build_pr_documentation`: doc-builder imported `SanaWMTransformer3DModel`
   from `diffusers.models.transformers` (not the diffusers top level) and
   that subpackage's `__init__.py` was missing the entry. Added the import.

* `doc-builder style src/diffusers docs/source --max_len 119` rewraps
  docstrings in the six SANA-WM files (transformer, kernels, pipeline,
  refiner, output, cam_utils) to the repo-wide 119-column limit. No
  behaviour change — purely whitespace inside docstrings.
* `make fix-copies` regenerates `dummy_pt_objects.py` and
  `dummy_torch_and_transformers_objects.py` to add `DummyObject` stubs
  for the three new public classes (`SanaWMTransformer3DModel`,
  `SanaWMPipeline`, `SanaWMLTX2Refiner`), so `from diffusers import …`
  gives the standard "missing backend" message on installs without
  torch / transformers.

Verified: `make quality` passes (ruff check, ruff format check,
doc-builder style check_only, check_doc_toc). Test suite still
15 passed / 1 skipped.

Drop 23 unused symbols (1451 lines) from `transformer_sana_wm.py` that
aren't reachable from the public SANA-WM checkpoint's
`SanaMSVideoCamCtrl` -> `SanaVideoMSCamCtrlBlock` -> GDN/UCPE attention
path. Each was verified to have zero call sites outside of its own
definition (or only within other-deleted items).

Classes:
* `SanaMS`, `SanaMSBlock` — alternative `Sana` subclass + block, not
  used by the SANA-WM checkpoint (which goes through `SanaMSVideoCamCtrl`).
* `ChunkCausalAttention`, `CachedCausalAttention`,
  `ChunkedLiteLAReLURope`, `LiteLAReLURope` — chunk-causal / cached
  attention variants and their common base; SANA-WM uses the bidi GDN
  path. `LiteLAReLURope` had only the three (now-deleted) subclasses
  referencing it.
* `PAGCFGIdentitySelfAttnProcessorLiteLA`,
  `PAGIdentitySelfAttnProcessorLiteLA`, `SelfAttnProcessorLiteLA`,
  `SelfAttnProcessorLiteLAReLURope` — PAG processors; we don't expose
  PAG in the SANA-WM pipeline.
* `ChunkGLUMBConvTemp`, `CachedGLUMBConvTemp`, `MBConvPreGLU` —
  alternative FFN/conv blocks; the checkpoint uses `GLUMBConvTemp`.
* `MaskFinalLayer`, `DecoderLayer` — alternative final layers; the
  checkpoint uses `T2IFinalLayer`.
* `LabelEmbedder`, `CaptionEmbedderDoubleBr` — alternative embedders;
  the checkpoint uses `CaptionEmbedder`.

Helpers:
* `set_grad_checkpoint`, `prepare_prompt_ar`, `resize_and_crop_tensor`,
  `generate_temporal_head_mask_mod`, `is_chunk_causal_request`,
  `get_chunk_index_from_config` — training-only or
  chunk-causal-mode helpers with no inference call sites.

Verified after the diff:
* `make quality` clean.
* CPU test suite 15 passed / 1 skipped (slow GPU integration).
* `import diffusers ; SanaWMPipeline / SanaWMTransformer3DModel /
  SanaWMLTX2Refiner` resolve identically.

Two PR-CI failures, both ours:

* `check_torch_dependencies`: line 33 hard-imported `transformers`, which
  isn't present in the minimum-deps environment. Move the lone
  `AutoModelForCausalLM` use site inside `initialize_gemma_params` (a
  training-only helper).
* `check_repository_consistency`: `SanaWMTransformer3DModel.forward`'s
  docstring was missing entries for `mask` and `return_dict`. Added.

Per @dg845's review comments on pipeline_sana_wm.py:

* Read the VAE spatial/temporal strides from the ``vae`` component
  (``vae_spatial_compression_ratio`` / ``vae_temporal_compression_ratio``)
  instead of hardcoding 32/8, mirroring LTX2Pipeline.
* Inline the stage-1 sampling loop into ``__call__`` and factor the noise
  init into a standard ``prepare_latents`` method.
* Use ``self.scheduler`` for the flow-matching Euler steps instead of
  constructing a new scheduler per call.
* Drive the sampling loop with ``self.progress_bar`` (respects
  ``set_progress_bar_config``) instead of a bare tqdm.
* Move input validation/normalization into a ``check_inputs`` method.
* Move first-frame resize+center-crop, [-1, 1] normalization and the
  intrinsics rescale into a ``SanaWMImageProcessor(VaeImageProcessor)``
  subclass (new ``image_processor.py``).
* Add ``generator`` as a ``__call__`` argument (``seed`` kept as a
  convenience shortcut).
* Post-process decoded latents with ``VideoProcessor.postprocess_video``
  rather than a hand-rolled conversion; drop the ``.cpu()`` casts on the
  ``output_type="latent"`` path.
* Move the pipeline README into the docs (docs/.../sana_wm.md); delete
  the in-package README.

Per @dg845's review: the refiner has components (transformer, connectors,
text encoder, tokenizer) and runs a denoising loop, so it fits better as a
pipeline than a ModelMixin.

* SanaWMLTX2Refiner now subclasses DiffusionPipeline and registers its
  components via ``register_modules`` (dropping the bespoke
  ``from_pretrained`` / ``save_pretrained``); standard load/save now handle
  the ``refiner/`` subfolder.
* Add a ``FlowMatchEulerDiscreteScheduler`` component (shift=1.0) and drive
  the Euler steps through ``scheduler.step`` / ``scheduler.scale_noise``
  (single-shot and per-AR-block) instead of hand-rolled updates. Numerically
  equivalent to the previous flow-matching update.
* Rename the entry point ``refine_latents`` -> ``__call__``; add a ``device``
  arg so the parent can hand it the execution device without a bulk move.
* SanaWMPipeline: keep the refiner as an optional nested component; free the
  parent's GPU weights before running it (it manages its own sub-module
  placement) and bring the VAE back for decode.
* Conversion script emits the new refiner layout (model_index.json +
  scheduler/ + tokenizer/); test asserts the refiner is a DiffusionPipeline
  with the canonical AR ``__call__`` defaults.

Validated end-to-end on 1xH100 (stage-1 + refiner on the official demo,
coherent video output).